import csv
import os
from typing import Any, Dict, List, Optional

from oneconvert.datasource.base import BaseDataSource


class CsvDataSource(BaseDataSource):
    def read(self, encoding='utf-8', mode='list', key_column: Optional[str] = None, **kwargs):
        """读取CSV文件并转成JSON数据
        :param encoding: 文件编码（默认utf-8）
        :param mode: 转换模式
                     - 'list': 转换为字典列表（默认）
                     - 'dict': 转换为以指定列或者指定行作为键的字典对象
        :param key_column: mode为'dict'时，指定作为键的列名
        :return JSON数据
        """
        if not os.path.isfile(self.uri):
            raise ValueError(f'the file {self.uri} not found')

        data: Any = []
        with open(self.uri, 'r', encoding=encoding) as f:
            reader = csv.DictReader(f)
            rows = list(reader)

            if mode == 'list':
                data = rows
            elif mode == 'dict':
                if not key_column:
                    raise ValueError("key_column required when mode is dict")
                if key_column not in reader.fieldnames:
                    raise ValueError(f"column '{key_column}' not found")
                data = {row[key_column]: row for row in rows}
            else:
                raise ValueError(f"unsupported mode '{mode}', should be in [list, dict]")

        return data

    def write(self, data: List | Dict,
              headers: List = None,
              encoding='utf-8',
              delimiter: str = ',',
              quotechar: str = '"',
              lineterminator: str = '\n', **kwargs):
        """保存json数据至csv
        :param data: 要写入的数据（字典列表，每个字典代表一行，键为列名）
        :param headers: 自定义表头列表，None 则使用 data 中第一个字典的键
        :param encoding: 文件编码（默认 utf-8）
        :param delimiter: 列分隔符（默认逗号 ,）
        :param quotechar: 字段引用符（默认双引号 "）
        :param lineterminator: 行终止符（默认 \n）
        """
        if data is None:
            data = []

        csv_data = []
        if isinstance(data, List):
            csv_data = data
        elif isinstance(data, Dict):
            csv_data = data.values()
        else:
            raise ValueError('data format should be List or Dict')

        if headers is None:
            headers = []
            if csv_data:
                headers = csv_data[0].keys()

        dir_name = os.path.dirname(self.uri)
        if dir_name and not os.path.exists(dir_name):
            os.makedirs(dir_name, exist_ok=True)

        with open(self.uri, 'w', encoding=encoding, newline='') as f:
            writer = csv.DictWriter(
                f,
                fieldnames=headers,
                delimiter=delimiter,
                quotechar=quotechar,
                lineterminator=lineterminator,
                quoting=csv.QUOTE_MINIMAL
            )

            writer.writeheader()
            writer.writerows(csv_data)
